package goutl

func NewSet[V comparable](cap int) *Set[V] {
	return &Set[V]{
		m: make(map[V]struct{}, cap),
	}
}

type Set[V comparable] struct {
	m map[V]struct{}
}

func (s *Set[V]) Add(v V) {
	s.m[v] = exist
}

func (s *Set[V]) Combine(set *Set[V]) {
	for v := range set.m {
		s.m[v] = exist
	}
}

func (s *Set[V]) Remove(v V) {
	delete(s.m, v)
}

func (s *Set[V]) Size() int {
	return len(s.m)
}

func (s *Set[V]) Empty() bool {
	return len(s.m) == 0
}

func (s *Set[V]) Contain(v V) bool {
	_, ok := s.m[v]
	return ok
}

func (s *Set[V]) Foreach(fn func(V)) {
	for v := range s.m {
		fn(v)
	}
}

func (s *Set[V]) ToArr() []V {
	arr := make([]V, 0, s.Size())

	s.Foreach(func(v V) {
		arr = append(arr, v)
	})

	return arr
}

func NewCustomSet[V any, K comparable](cap int, keyFun func(V) K) *CustomSet[V, K] {
	return &CustomSet[V, K]{
		m:   make(map[K]V, cap),
		key: keyFun,
	}
}

type CustomSet[V any, K comparable] struct {
	m   map[K]V
	key func(V) K
}

func (s *CustomSet[V, K]) Add(v V) {
	s.m[s.key(v)] = v
}

func (s *CustomSet[V, K]) Remove(v V) {
	delete(s.m, s.key(v))
}

func (s *CustomSet[V, K]) RemoveKey(k K) {
	delete(s.m, k)
}

func (s *CustomSet[V, K]) Size() int {
	return len(s.m)
}

func (s *CustomSet[V, K]) Contain(v V) bool {
	_, ok := s.m[s.key(v)]
	return ok
}

func (s *CustomSet[V, K]) ContainKey(k K) bool {
	_, ok := s.m[k]
	return ok
}

func (s *CustomSet[V, K]) Foreach(fn func(V)) {
	for _, v := range s.m {
		fn(v)
	}
}

func (s *CustomSet[V, K]) ToArr() []V {
	arr := make([]V, 0, s.Size())

	s.Foreach(func(v V) {
		arr = append(arr, v)
	})

	return arr
}

func (s *CustomSet[V, K]) GetCustomKeys() []K {
	keys := make([]K, 0, s.Size())

	s.Foreach(func(v V) {
		keys = append(keys, s.key(v))
	})

	return keys
}
